package com.wuwangfu.window.aggre;

import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.datastream.WindowedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

import java.time.Duration;

/**
 * @Author: jcshen
 * @Date: 2023-03-09
 * @PackageName: com.wuwangfu.window.aggregate
 * @ClassName: WindowAggregateFunction
 * @Description: 累加当前窗口的数据，并与历史数据进行累加
 * @Version: 1.0.0
 * <p>
 * https://nightlies.apache.org/flink/flink-docs-release-1.14/docs/dev/datastream/operators/windows/#aggregatefunction
 */
public class WindowAggregateFunction {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.enableCheckpointing(10000);

//        1000,spark,1
//        1001,spark,1
//        5000,spark,2
//        6000,hive,2
//        9000,spark,3
//        10000,hadoop,2
        DataStreamSource<String> lines = env.socketTextStream("localhost", 8888);
        //生成水位线
        SingleOutputStreamOperator<String> dataWatermark = lines.assignTimestampsAndWatermarks(WatermarkStrategy
                //生成水位线策略
                .<String>forBoundedOutOfOrderness(Duration.ZERO)
                //提取水位线时间
                .withTimestampAssigner((line, timestamp) -> Long.parseLong(line.split(",")[0]))

        );
        //
        SingleOutputStreamOperator<Tuple2<String, Integer>> maped = dataWatermark.map(new MapFunction<String, Tuple2<String, Integer>>() {
            @Override
            public Tuple2<String, Integer> map(String value) throws Exception {
                String[] fields = value.split(",");
                return Tuple2.of(fields[1], Integer.parseInt(fields[2]));
            }
        });
        //分组
        KeyedStream<Tuple2<String, Integer>, String> keyed = maped.keyBy(t -> t.f0);
        //开窗
        WindowedStream<Tuple2<String, Integer>, String, TimeWindow> windowed = keyed.window(TumblingEventTimeWindows.of(Time.seconds(5)));
        //窗口内聚合
        SingleOutputStreamOperator<Tuple2<String, Integer>> aggregate = windowed.aggregate(new MyAggreFunc(), new MyWindowFunc());

        aggregate.print();


        env.execute();
    }

    private static class MyAggreFunc implements AggregateFunction<Tuple2<String, Integer>, Integer, Integer> {
        /**
         * 创建一个初始值
         *
         * @return
         */
        @Override
        public Integer createAccumulator() {
            return 0;
        }

        /**
         * 输入一条数据，与初始值或中间累加的结果进行聚合
         *
         * @param value
         * @param accumulator
         * @return
         */
        @Override
        public Integer add(Tuple2<String, Integer> value, Integer accumulator) {
            return value.f1 + accumulator;
        }

        /**
         * 返回结果
         *
         * @param accumulator
         * @return
         */
        @Override
        public Integer getResult(Integer accumulator) {
            return accumulator;
        }


        /**
         * 如果使用的是非SessionWindow，可以不实现
         *
         * @param a
         * @param b
         * @return
         */
        @Override
        public Integer merge(Integer a, Integer b) {
            return null;
        }
    }

    private static class MyWindowFunc extends ProcessWindowFunction<Integer, Tuple2<String, Integer>, String, TimeWindow> {

        private transient ValueState<Integer> valueState;

        @Override
        public void open(Configuration parameters) throws Exception {

            ValueStateDescriptor<Integer> valueDesc = new ValueStateDescriptor<>("wc-state", Integer.class);
            valueState = getRuntimeContext().getState(valueDesc);
        }

        @Override
        public void process(String key, Context context, Iterable<Integer> iterable, Collector<Tuple2<String, Integer>> out) throws Exception {
            Integer state = valueState.value();
            if (state == null) {
                state = 0;
            }
            //获取窗口聚合后输出的结果
            Integer value = iterable.iterator().next();
            int acc = value + state;
            //更新状态
            valueState.update(acc);
            //输出聚合结果
            out.collect(Tuple2.of(key,acc));
        }
    }
}
